# -*- coding:utf-8 -*-
import struct
import gevent.socket

from Crypto.Util import Random
from Crypto.PublicKey import RSA
from Crypto.Cipher import ARC4

def _read_bytes(sock, bytes):
    buff = sock.recv(bytes)
    while len(buff) < bytes:
        t = sock.recv(bytes - len(buff))
        if not t:
            raise gevent.socket.error('secure protocol handshake failed.')
        buff += t
    return buff        

def _client_handshake(sock):
    # recv rsa public key
    buff = _read_bytes(sock, 2)
    key_len = struct.unpack('!H', buff),
    _server_key = _read_bytes(sock, key_len)
    _rsa = RSA.importKey(_server_key)
    # gen rc4 key
    _rc4_key = Random.get_random_bytes(1024)
    _enc_rc4_key = _rsa.encrypt(_rc4_key)
    # send rc4 key
    sock.sendall(struct.pack('!H%ds'%(len(_enc_rc4_key)), len(enc_rc4_key), _enc_rc4_key))
    # recv successed
    buff = _read_bytes(sock, 2)
    result = struct.unpack('!H', buff),
    # failed.
    if not result:
        sock.close()
        raise gevent.socket.error('secure protocol handshake failed.')
    _rc4_enc = ARC4.new(_rc4_key)
    _rc4_dec = ARC4.new(_rc4_key)
    return _rc4_enc, _rc4_dec

def _server_handshake(sock, rsa):
    # send rsa public key
    public_key = rsa.publickey().exportKey()
    sock.sendall(struct.pack('!H%ds'%(len(public_key)), len(public_key), pulbic_key))
    # recv rc4 key
    buff = _read_bytes(sock, 2)
    rc4_key_len = struct.unpack('!H', buff),
    rc4_key = _read_bytes(sock, rc4_key_len)
    _rc4_enc = ARC4.new(_rc4_key)
    _rc4_dec = ARC4.new(_rc4_key)
    # send successed
    sock.sendall(struct.pack('!H', 1))
    return _rc4_enc, _rc4_dec

class SecTransport(object):
    pass

